import math
from class_decomposition import *

def decompose(FR2, eps, eps_cutoff = 0, disp = False):
	'''
	decomposition in elasticity form, at n: 2->3
	'''
	# compute optimal n, under discrete choice
	if disp: 
		print("-----------------------------------------------------")
		print("(1). Compute optimal fertility and evaluate mother type")
	
	FR2.evalDiscreteN(disp = False)

	FR2.evalQQ_DiscreteN()

	if FR2.n_opt[0] == 2:
		type2to3 = 'A'
	elif FR2.n_opt[0] == 3 and eps < eps_cutoff:
		type2to3 = 'B'
	else: 
		type2to3 = 'C'
 
	# store type characteristics
	temp_type = pd.DataFrame(
		[{
			"type2to3": type2to3,
		}]
	)

	if disp:
		print("FamilyRation() class information: ")
		print("n_opt:", FR2.n_opt[0] )
		print("type2to3:", type2to3)
		print("q_opt", FR2.q_opt)

	# skip decomposition if both types are "C"
	if type2to3 == "C" :

		if disp: 
			print("(2). Skip type C (including n_opt = 1)")
			print("-----------------------------------------------------")
		return temp_type

	elif type2to3 in ["A", "B"]:

		# create an instance of decomposition
		D = Decomposition(FR2)
		
		# decompose at n = 2.5
		if disp:
			print("(2). Decompose at n == 2.5, return three effects")

		D.nxset(n = [2.5], x0 = [0.1, 0.1])
		D.run()
		D.getDataDecomp_e()

		# define the list of column names to extract from D.dData_e
		extract_columns = ['q', 'total', 'price', 'substitution', 'income']
		temp2half = D.dData_e[extract_columns]
		temp2half.columns = list( map(lambda x: x+"_2half", extract_columns) )


		if disp:
			print("-----------------------------------------------------")

		return pd.concat( [temp_type, temp2half], axis = 1 )


